from vanna.flask import VannaFlaskApp
from vanna.chromadb import ChromaDB_VectorStore
import threading
from APP.config.conf import settings
from APP.common.middleware.xinference_api.xinferenceEmbeddings import Xinference_Embeddings_Function
from APP.common.middleware.xinference_api.xinference import Xinference
import asyncio

local_ip = settings.VANNA_LOCAL_IP 
api_key = 'd7379202cee64ec7bd499e62e77d3a90.CHKkv8oYO72IEsp5'
xinference_embedding_function = Xinference_Embeddings_Function(
    config={"base_url":settings.VANNA_BASE_URL,"api_key": '',"model_uid":settings.VANNA_EMBEDDING_MODEL_UID}
)

class CustomFlaskApp(VannaFlaskApp):
    def run(self, port):
        print("Your app is running at:")
        print(f"{local_ip}:{port}")
        self.flask_app.run(host="0.0.0.0", port=port, use_reloader=False, debug=True)

class init_Vanna(ChromaDB_VectorStore, Xinference):
    def __init__(self, config=None, vsconfig=None):
        ChromaDB_VectorStore.__init__(self, config=vsconfig)
        Xinference.__init__(self, config=config)


class userVanna:
    def __init__(self, init_Vanna=init_Vanna, customflaskapp=CustomFlaskApp):
        self.vsconfig = self.get_or_create_vsconfig_path()
        self.user_Vanna = init_Vanna(config={"base_url":settings.VANNA_BASE_URL,"api_key": '', "llm_model_uid": settings.VANNA_LLM_MODEL_UID}, vsconfig=self.vsconfig)
        self.port = None
        self.customflaskapp = customflaskapp
        self.port_event = threading.Event()
    
    def get_or_create_vsconfig_path(self):
        
        return {
            "embedding_function": xinference_embedding_function,
            "n_results_sql": 5,
            "n_results_documentation": 5,
            "n_results_ddl": 3,
            "path": './embedding_db/101916446.chroma'
        }

    def web_server(self):
        app = self.customflaskapp(self.user_Vanna, allow_llm_to_see_data=True)
        self.port = self.find_free_port()
        self.port_event.set()
        app.run(port=self.port)

    def find_free_port(self):
        import socket
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(('', 0))
            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            return s.getsockname()[1]

    def start_web_server(self):
        thread = threading.Thread(target=self.web_server)
        thread.start()
        self.port_event.wait()
        return self.port
    
    async def ask(self, question, df, sql):
        # 使用run_in_executor将同步IO操作放到线程池中执行
        loop = asyncio.get_event_loop()
        
        # 异步生成plotly代码
        fig_code = await loop.run_in_executor(
            None, 
            self.user_Vanna.generate_plotly_code,
            question, sql, df
        )

        if fig_code:
            # 异步获取plotly图表
            fig = await loop.run_in_executor(
                None,
                self.user_Vanna.get_plotly_figure,
                fig_code, df
            )
            # fig.show()

            fig_str = fig.to_json()
            return fig_str
        return None